import cPickle as pickle
from optparse import OptionParser

import tables
from scipy.stats.stats import pearsonr 

from mp3row import MP3Row
from matchrow import MatchRow


# Match algorithm:
# 1. Copy song md5sums into queue.
# 2. Grab a test song from md5 queue.
# 3. Select row from Generate ntuple corresponding to the test song's md5.
# 4. Select all rows from the Generate ntuple with length within 10s of that of
#    the test song.
# 5. For each row in the selection, compare WAV data.
# 6. If a row is strongly correllated with the test song, one of the following
#    must be true:
#      a. Neither the test song nor the correllated song are in a set in the match list.
#         i.   If so, create a set with both the test and correllated songs, and
#              append that set to the match list.
#         ii.  Remove the correllated song from the md5 queue.
#         iii. Continue from (5)
#      b. The test song is in a set in the match list, but the correllated song is not.
#         i.   If so, add the correllated song to the set in the match list 
#              which includes the test song.
#         ii.  Remove the correllated song from the md5 queue.
#         iii. Continue from (5)
#      c. The correllated song is in a set in the match list, but the test song is not.
#         i.   If so, add the test song to the set in the match list 
#              which includes the correllated song.
#         iii. Continue from (5)
# 7. Remove the test song from the md5 queue.
# 8. Continue from (2)

def match():
    op = OptionParser("mp3same match [opts] file1.mp3 | dir1 [file2.mp3 | dir2 [...]]")
    op.add_option('-o', '--output', dest='output', default='mp3same_match.h5',
                  help='Redirect output to PATH', metavar='PATH')
    op.add_option('-r', '--r-cut', dest='rcut', type='float', default=0.8,
                  help='Consider two songs as duplicates when R value greater than RCUT', 
                  metavar='RCUT')

    (opts,args) = op.parse_args()

    if len(args) < 2:
        sys.stderr.write("`mp3same.py match' requires at least one argument.\n")
        sys.stderr.write("See `mp3same.py match -h' for more.\n")
        sys.stderr.flush()
        return 1

    infp    = tables.openFile(args[1])
    ofp     = tables.openFile(opts.output + '.h5', 'w', 'mp3same match output file')
    pkl_in  = open(args[1][:-2] + 'pkl', 'r')
    out_table = ofp.createTable(ofp.root, "mnt", MatchRow, 'mp3same match data') 

    (md5dict, trash) = pickle.load(pkl_in)
#    md5dict = set(infp.root.hdrnt[:]['mp3md5'])
    matches = list()

    nfiles = 1; 
    for testrow in infp.root.hdrnt:
        print nfiles, 'of', len(infp.root.hdrnt)
        nfiles += 1
  
        if testrow['mp3md5'] not in md5dict:
            continue

        if testrow['fname'] in trash:
            continue

        if testrow['nchannels'] == 1:
            sys.stderr.write("`" + testrow['fname'] + "'" 
                             + " is a single-channel track--skipping correllation test...\n")
            sys.stderr.flush()
            continue
        
        selstr = '(length < ' + str(testrow['length'] + 10)
        selstr += ') & (length > ' + str(testrow['length'] - 10) + ')'
        selstr += " & (mp3md5 != '" + testrow['mp3md5'] + "')" 

        test_data = infp.getNode(infp.root.data, testrow['arrname']).read()
        
        for selrow in infp.root.hdrnt.where(selstr):
            seldata = infp.getNode(infp.root.data, selrow['arrname']).read()
            len_to_read = min((len(test_data), len(seldata)))
            
            (rval, pval) = pearsonr(test_data[:len_to_read], 
                                    seldata[:len_to_read])

            out_table.row['md5a'] = testrow['mp3md5']
            out_table.row['md5b'] = selrow['mp3md5']
            out_table.row['r']    = rval
            out_table.row['p']    = pval

            out_table.row.append()
            out_table.flush()

            if rval < opts.rcut:
                continue

            match_exists = False
            for matchdict in matches:
                if testrow['fname'] in matchdict and selrow['fname'] in matchdict:
                    sys.stderr.write("`" + testrow['fname'] + "' and `" + selrow['fname']
                                     + "' already in a matchdict! WTF!\n")
                    sys.stderr.flush()
                    break;
                elif testrow['fname'] in matchdict:
                    matchdict[selrow['fname']] = selrow['mp3size']
                    del md5dict[selrow['mp3md5']]
                    match_exists = True
                elif selrow['fname'] in matchdict:
                    matchdict[testrow['fname']] = selrow['mp3size']
                    match_exists = True
            
            
            if not match_exists:
                matches.append({testrow['fname'] : testrow['mp3size'], selrow['fname']: selrow['mp3size']})
                del md5dict[selrow['mp3md5']]


        del md5dict[testrow['mp3md5']]
           
   
    infp.close()
    ofp.close()
    
    pfp = open(opts.output + '.pkl', 'wb')
    pickle.dump(matches, pfp)
    pfp.close()

    
    
